{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define new problem type and data reading function\n",
    "\n",
    "We'll use IMDB dataset as example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/data3/yjp/bert-multitask-learning\n"
     ]
    }
   ],
   "source": [
    "cd ../"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bert_multitask_learning import (get_or_make_label_encoder, FullTokenizer, \n",
    "                                     create_single_problem_generator, train_bert_multitask, \n",
    "                                     eval_bert_multitask, DynamicBatchSizeParams, TRAIN, EVAL, PREDICT, preprocessing_fn)\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_problem_type = {'imdb_cls': 'cls'}\n",
    "\n",
    "@preprocessing_fn\n",
    "def imdb_cls(params, mode):\n",
    "\n",
    "    # get data\n",
    "    (train_data, train_labels), (test_data, test_labels) = keras.datasets.imdb.load_data(num_words=10000)\n",
    "    label_encoder = get_or_make_label_encoder(params, 'imdb_cls', mode, train_labels+test_labels)\n",
    "    word_to_id = keras.datasets.imdb.get_word_index()\n",
    "    index_from=3\n",
    "    word_to_id = {k:(v+index_from) for k,v in word_to_id.items()}\n",
    "    word_to_id[\"<PAD>\"] = 0\n",
    "    word_to_id[\"<START>\"] = 1\n",
    "    word_to_id[\"<UNK>\"] = 2\n",
    "    id_to_word = {value:key for key,value in word_to_id.items()}\n",
    "\n",
    "    train_data = [[id_to_word[i] for i in sentence] for sentence in train_data]\n",
    "    test_data = [[id_to_word[i] for i in sentence] for sentence in test_data]\n",
    "    \n",
    "    if mode == TRAIN:\n",
    "        input_list = train_data\n",
    "        target_list = train_labels\n",
    "    else:\n",
    "        input_list = test_data\n",
    "        target_list = test_labels\n",
    "    \n",
    "    return input_list, target_list\n",
    "new_problem_process_fn_dict = {'imdb_cls': imdb_cls}\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Model\n",
    "\n",
    "Please make sure you're using the correct checkpoint to initialize model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding new problem imdb_cls, problem type: cls\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:3\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:3\n",
      "INFO:tensorflow:Configured nccl all-reduce.\n",
      "INFO:tensorflow:Initializing RunConfig with distribution strategies.\n",
      "INFO:tensorflow:Not using Distribute Coordinator.\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'models/imdb_cls_ckpt', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
      "graph_options {\n",
      "  rewrite_options {\n",
      "    meta_optimizer_iterations: ONE\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7fceaec24240>, '_device_fn': None, '_protocol': None, '_eval_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7fceaec24240>, '_experimental_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7fcf795f3748>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None}\n",
      "INFO:tensorflow:Create RestoreCheckpointHook.\n",
      "INFO:tensorflow:Skipping training since max_steps has already saved.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow_estimator.python.estimator.estimator.Estimator at 0x7fceaec24a20>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "params = DynamicBatchSizeParams()\n",
    "params.init_checkpoint = 'models/cased_L-12_H-768_A-12'\n",
    "tf.logging.set_verbosity(tf.logging.DEBUG)\n",
    "train_bert_multitask(problem='imdb_cls', num_gpus=1, \n",
    "                     num_epochs=10, params=params, \n",
    "                     problem_type_dict=new_problem_type, processing_fn_dict=new_problem_process_fn_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate Model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Params problem assigned. Problem list: ['imdb_cls']\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:3\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:3\n",
      "INFO:tensorflow:Configured nccl all-reduce.\n",
      "INFO:tensorflow:Initializing RunConfig with distribution strategies.\n",
      "INFO:tensorflow:Not using Distribute Coordinator.\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'models/imdb_cls_ckpt', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
      "graph_options {\n",
      "  rewrite_options {\n",
      "    meta_optimizer_iterations: ONE\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7fcea60f1080>, '_device_fn': None, '_protocol': None, '_eval_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7fcea60f1080>, '_experimental_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7fceaec27550>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None}\n",
      "WARNING:tensorflow:From /data3/yjp/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n",
      "WARNING:tensorflow:From /data3/yjp/anaconda3/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py:429: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "tf.py_func is deprecated in TF V2. Instead, use\n",
      "    tf.py_function, which takes a python function which manipulates tf eager\n",
      "    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n",
      "    an ndarray (just call tensor.numpy()) but having access to eager tensors\n",
      "    means `tf.py_function`s can use accelerators such as GPUs as well as\n",
      "    being differentiable using a gradient tape.\n",
      "    \n",
      "INFO:tensorflow:Calling model_fn.\n",
      "WARNING:tensorflow:From /data3/yjp/bert-multitask-learning/bert_multitask_learning/bert/modeling.py:673: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dense instead.\n",
      "DEBUG:tensorflow:Converted call: <function stop_grad at 0x7fceaef32e18>; owner: None\n",
      "DEBUG:tensorflow:Converting <function stop_grad at 0x7fceaef32e18>\n",
      "DEBUG:tensorflow:Compiled output of <function stop_grad at 0x7fceaef32e18>:\n",
      "\n",
      "def stop_grad(global_step, tensor, freeze_step):\n",
      "  try:\n",
      "    with ag__.function_scope('stop_grad'):\n",
      "      cond_1 = ag__.gt(freeze_step, 0)\n",
      "\n",
      "      def if_true_1():\n",
      "        with ag__.function_scope('if_true_1'):\n",
      "          tensor_2, = tensor,\n",
      "          cond = ag__.lt_e(global_step, freeze_step)\n",
      "\n",
      "          def if_true():\n",
      "            with ag__.function_scope('if_true'):\n",
      "              tensor_1, = tensor_2,\n",
      "              tensor_1 = tf.stop_gradient(tensor_1)\n",
      "              return tensor_1\n",
      "\n",
      "          def if_false():\n",
      "            with ag__.function_scope('if_false'):\n",
      "              return tensor_2\n",
      "          tensor_2 = ag__.if_stmt(cond, if_true, if_false)\n",
      "          return tensor_2\n",
      "\n",
      "      def if_false_1():\n",
      "        with ag__.function_scope('if_false_1'):\n",
      "          return tensor\n",
      "      tensor = ag__.if_stmt(cond_1, if_true_1, if_false_1)\n",
      "      return tensor\n",
      "  except:\n",
      "    ag__.rewrite_graph_construction_error(ag_source_map__)\n",
      "\n",
      "\n",
      "\n",
      "stop_grad.autograph_info__ = {}\n",
      "\n",
      "\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "WARNING:tensorflow:From /data3/yjp/anaconda3/lib/python3.7/site-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use standard file APIs to check for files with this prefix.\n",
      "INFO:tensorflow:Restoring parameters from models/imdb_cls_ckpt/model.ckpt-7812\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing Inputs: 100%|██████████| 25000/25000 [04:57<00:00, 83.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'imdb_cls_Accuracy': 0.91272, 'imdb_cls_Accuracy Per Sequence': 0.91272}\n"
     ]
    }
   ],
   "source": [
    "print(eval_bert_multitask(problem='imdb_cls', num_gpus=1, \n",
    "                     params=params, eval_scheme='acc',\n",
    "                     problem_type_dict=new_problem_type, processing_fn_dict=new_problem_process_fn_dict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
